import numpy as np
import cv2


def sort_quad_vertices(vertices):
    """
    功能：将四边形顶点按「左上角→右上角→右下角→左下角」顺序排序
    输入：vertices - 多边形逼近得到的顶点数组（shape=(4,1,2)）
    输出：sorted_pts - 排序后的顶点数组（shape=(4,2)，float32）
    """
    # 去除冗余维度，转为4行2列的坐标数组
    pts = vertices.reshape((4, 2)).astype(np.float32)

    # 1. 按x+y之和排序：左上角（sum最小）、右下角（sum最大）
    sum_xy = pts.sum(axis=1)
    tl_idx = np.argmin(sum_xy)  # 左上角索引
    br_idx = np.argmax(sum_xy)  # 右下角索引

    # 2. 排除已确定的角点，剩余两个为右上角和左下角
    remaining_idx = [i for i in range(4) if i not in [tl_idx, br_idx]]
    tr_idx, bl_idx = remaining_idx

    # 3. 区分右上角和左下角：x坐标大的是右上角
    if pts[tr_idx, 0] < pts[bl_idx, 0]:
        tr_idx, bl_idx = bl_idx, tr_idx

    # 按顺序返回四个角点
    sorted_pts = np.array([pts[tl_idx], pts[tr_idx], pts[br_idx], pts[bl_idx]], dtype=np.float32)
    return sorted_pts


def get_target_dimensions(sorted_pts):
    """
    功能：根据排序后的顶点计算目标图像真实宽高（避免拉伸）
    输入：sorted_pts - 排序后的顶点数组（shape=(4,2)）
    输出：target_w, target_h - 目标宽高（整数）
    """
    tl, tr, br, bl = sorted_pts

    # 计算水平方向宽度（取上下边最大值，确保完整覆盖）
    width_top = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    width_bottom = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    target_w = int(np.ceil(max(width_top, width_bottom)))  # 向上取整避免裁剪

    # 计算垂直方向高度（取左右边最大值）
    height_left = np.sqrt(((bl[0] - tl[0]) ** 2) + ((bl[1] - tl[1]) ** 2))
    height_right = np.sqrt(((br[0] - tr[0]) ** 2) + ((br[1] - tr[1]) ** 2))
    target_h = int(np.ceil(max(height_left, height_right)))

    return target_w, target_h


def preprocess_image(img):
    """
    功能：图像预处理（灰度化→滤波→二值化），增强轮廓检测效果
    输入：img - 原始彩色图像
    输出：img_binary - 二值化图像（便于轮廓检测）
    """
    # 1. 灰度化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 2. 高斯滤波（去噪，核大小3x3更精细，sigmaX=0自动计算）
    blurred = cv2.GaussianBlur(gray, ksize=(3, 3), sigmaX=0)

    # 3. 自适应二值化（适应光影变化，比固定阈值更鲁棒）
    binary = cv2.adaptiveThreshold(
        blurred, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,  # 高斯加权自适应
        cv2.THRESH_BINARY_INV,  # 反色（目标为白色，背景为黑色）
        blockSize=11,  # 邻域块大小（奇数）
        C=2  # 阈值调整量（减去C降低误检）
    )

    return binary


# -------------------------- 主程序入口 --------------------------
if __name__ == "__main__":
    # 1. 读取图像（请确保路径正确，支持相对/绝对路径）
    img_path = "../images/test_wrap.png"
    img = cv2.imread(img_path)
    if img is None:
        print(f"错误：无法读取图像文件，路径：{img_path}")
        print("请检查：1.路径是否正确 2.文件是否存在 3.文件格式是否支持（jpg/png）")
        exit()

    # 2. 图像预处理
    img_binary = preprocess_image(img)

    # 3. 轮廓检测（仅保留最外层轮廓，减少冗余）
    contours, _ = cv2.findContours(
        img_binary.copy(),  # 传入副本避免修改原二值图
        cv2.RETR_EXTERNAL,  # 只检测最外层轮廓
        cv2.CHAIN_APPROX_SIMPLE  # 压缩轮廓点（减少计算量）
    )

    # 4. 筛选最大轮廓（目标卡片的轮廓）
    if len(contours) == 0:
        print("错误：未检测到任何轮廓，请检查图像对比度或预处理参数")
        cv2.imshow("Preprocessed Binary Image", img_binary)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        exit()

    # 按轮廓面积排序，取最大的一个
    max_contour = max(contours, key=lambda c: cv2.contourArea(c))
    max_area = cv2.contourArea(max_contour)
    print(f"检测到最大轮廓：面积={max_area:.2f}px²，边界点数量={len(max_contour)}")

    # 5. 多边形逼近（获取卡片的四个角点）
    perimeter = cv2.arcLength(max_contour, closed=True)  # 计算轮廓周长（闭合轮廓）
    # epsilon=0.015*perimeter：平衡精度与简化，确保逼近为四边形
    approx_vertices = cv2.approxPolyDP(max_contour, epsilon=0.015 * perimeter, closed=True)

    # 校验：必须是四边形才能进行透视矫正
    if len(approx_vertices) != 4:
        print(f"错误：多边形逼近结果为{len(approx_vertices)}个顶点，需为四边形")
        print("建议调整：1.减小epsilon系数（如0.01） 2.优化图像预处理（降低噪声）")
        # 显示轮廓供调试
        img_contour = img.copy()
        cv2.drawContours(img_contour, [max_contour], -1, (0, 0, 255), 2)
        cv2.imshow("Max Contour", img_contour)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        exit()

    # 6. 顶点排序（关键：确保与目标顶点顺序对应）
    sorted_src_pts = sort_quad_vertices(approx_vertices)

    # 7. 计算目标图像尺寸
    target_w, target_h = get_target_dimensions(sorted_src_pts)
    print(f"矫正后目标尺寸：宽={target_w}px，高={target_h}px")

    # 8. 定义目标顶点（与源顶点顺序严格一致：左上→右上→右下→左下）
    dst_pts = np.float32([
        [0, 0],  # 目标左上
        [target_w - 1, 0],  # 目标右上（减1避免坐标越界）
        [target_w - 1, target_h - 1],  # 目标右下
        [0, target_h - 1]  # 目标左下
    ])

    # 9. 透视变换（核心步骤）
    # 计算透视变换矩阵
    M = cv2.getPerspectiveTransform(sorted_src_pts, dst_pts)
    # 执行透视变换（输出尺寸为目标尺寸，无黑边）
    img_corrected = cv2.warpPerspective(img, M, (target_w, target_h))

    # 10. 可视化结果（标记关键信息，便于验证）
    img_visual = img.copy()
    # 绘制最大轮廓（红色）
    cv2.drawContours(img_visual, [max_contour], -1, (0, 0, 255), 2)
    # 绘制排序后的顶点（蓝色实心圆+编号）
    for i, (x, y) in enumerate(sorted_src_pts, 1):
        cv2.circle(img_visual, (int(x), int(y)), 5, (255, 0, 0), -1)  # 实心蓝点
        cv2.putText(
            img_visual, str(i), (int(x) + 8, int(y) + 8),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2
        )
    # 绘制外接矩形（青色）
    x_min = int(min(sorted_src_pts[:, 0]))
    y_min = int(min(sorted_src_pts[:, 1]))
    x_max = int(max(sorted_src_pts[:, 0]))
    y_max = int(max(sorted_src_pts[:, 1]))
    cv2.rectangle(img_visual, (x_min, y_min), (x_max, y_max), (255, 255, 0), 2)

    # 11. 显示与保存
    # 调整窗口大小（适配屏幕）
    cv2.namedWindow("Original with Markers", cv2.WINDOW_NORMAL)
    cv2.namedWindow("Corrected Image", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Original with Markers", 600, 800)
    cv2.resizeWindow("Corrected Image", 600, 800)

    # 显示图像
    cv2.imshow("Original with Markers", img_visual)
    cv2.imshow("Corrected Image", img_corrected)

    # 保存结果（可选，取消注释即可）
    # cv2.imwrite("../images/original_markers.png", img_visual)
    # cv2.imwrite("../images/corrected_card.png", img_corrected)
    # print("结果已保存至 ../images 目录")

    # 等待按键关闭窗口
    print("提示：按任意键关闭窗口")
    cv2.waitKey(0)
    cv2.destroyAllWindows()
